# -*- encoding: utf-8 -*-
'''
@File    :   simulation.py
@Time    :   2021/06/15 10:07:13
@Author  :   Fei gao 
@Contact :   feig@mail.bnu.edu.cn
BNU, Beijing, China
'''
import os
import subprocess
import numpy as np

def SIMU(diff, graph, node_prob, edge_prob, seeds, path, num_simulation=10000, cpu_core=20):
    """C implementation of Simulation

    Args:
        diff (str]): "SIR" or "SIS"
        graph (nx.DiGraph): the graph structure
        node_prob (dict): dict, map from node to nodes' probability
        edge_prob (dict): map from edges to edges' probability
        seeds (list): list of seeds
        num_simulation (int): [description]
        cpu_core (int): [description]

    Raises:
        RuntimeError: [description]

    Returns:
        [type]: [description]
    """
    # write txt file for simulation
    txt = open(".tmp_simulation.txt", "w")
    N = graph.number_of_nodes()
    E = graph.number_of_edges()
    S = len(seeds)
    txt.write("{} {} {}\n\n\n".format(N, E, S))
    seeds = [str(s) for s in seeds]
    txt.write(" ".join(seeds)+"\n\n\n")
    
    for n, p in node_prob.items():
        txt.write("{} {}\n".format(n, p))
    txt.write("\n\n")

    for e, p in edge_prob.items():
        txt.write("{} {} {}\n".format(e[0], e[1], p))
    txt.close()
    
    # run simulation
    cmd = "./src/simulation -dataset .tmp_simulation.txt -model {} -mc {} -maxSimSteps 50 -cpuCores {}".format(diff, num_simulation, cpu_core)
    res = subprocess.getstatusoutput(cmd)
    if int(res[0]) != 0:
        print(res[1])
        raise RuntimeError
    
    # read simulation results
    with open("spreadProcess.txt", "r") as f:
        spread = f.readlines()
    T, N, K = [int(x) for x in spread[0].strip().split(" ")]
    spreading = np.zeros((T*N, K))
    cnt = 0
    for line in spread[1:]:
        if len(line.strip()) == 0:
            continue
        for k, x in enumerate(line.strip().split(" ")):
            spreading[cnt, k] = float(x)
        cnt += 1
    assert cnt == T*N
    spreading = spreading.reshape(T, N, K)

    os.remove(".tmp_simulation.txt")
    os.remove("spreadProcess.txt")
    return spreading
